{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Grok example with live search"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary and Best Practices\n",
    "\n",
    "This notebook demonstrated successful Grok integration with AG2. Here are the key takeaways:\n",
    "\n",
    "### ✅ What Works\n",
    "\n",
    "1. **OpenAI Compatibility**: Grok works seamlessly with AG2's OpenAI client\n",
    "2. **Real-time Search**: The `extra_body` parameter enables Grok's search capabilities\n",
    "3. **Function Calling**: Proper function registration\n",
    "4. **Standard AG2 Patterns**: All existing AG2 agent patterns work with Grok\n",
    "\n",
    "\n",
    "### 📚 References\n",
    "\n",
    "- [x.ai Function Calling Guide](https://docs.x.ai/docs/guides/function-calling)\n",
    "- [AG2 Documentation](https://ag2.ai/)\n",
    "- [Grok API Documentation](https://docs.x.ai/)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Before running this notebook, ensure you have:\n",
    "\n",
    "1. **API Key**: Get your Grok API key from [x.ai](https://x.ai/)\n",
    "2. **Environment Variable**: \n",
    "   ```bash\n",
    "   export XAI_API_KEY=\"your-actual-grok-api-key\"\n",
    "   ```\n",
    "3. **Dependencies**:\n",
    "   ```bash\n",
    "   pip install ag2[openai]\n",
    "   ```\n",
    "\n",
    "**Security Note**: Never hardcode API keys in notebooks. Always use environment variables for authentication."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2: Grok Function Calling (with Tools)\n",
    "This example demonstrates how to use Grok's function calling capabilities with AG2.\n",
    "**Features:**\n",
    "- Register Python functions as tools for Grok to call\n",
    "- Use OpenAI-compatible function calling interface\n",
    "- Best practice: Pass functions to the agent, *not* to `LLMConfig`\n",
    "```python\n",
    "from typing import Annotated\n",
    "def get_weather(city: Annotated[str, \"The city name\"]) -> str:\n",
    "    \"\"\"Get current weather for a city.\"\"\"\n",
    "    # This is a mock function - in reality, you'd call a weather API\n",
    "    return f\"The current weather in {city} is sunny with a temperature of 22°C.\"\n",
    "def calculate_math(expression: Annotated[str, \"Mathematical expression to evaluate\"]) -> str:\n",
    "    \"\"\"Calculate a mathematical expression safely.\"\"\"\n",
    "    try:\n",
    "        # Simple evaluation - in production, use a safer math parser\n",
    "        result = eval(expression.replace(\"^\", \"**\"))\n",
    "        return f\"The result of {expression} is {result}\"\n",
    "    except Exception:\n",
    "        return f\"Could not evaluate the expression: {expression}\"\n",
    "# Configure Grok for function calling (do NOT put tools in LLMConfig)\n",
    "function_config = LLMConfig(\n",
    "    {\n",
    "        \"model\": \"grok-4\",\n",
    "        \"api_key\": os.getenv(\"XAI_API_KEY\"),\n",
    "        \"base_url\": \"https://api.x.ai/v1\",\n",
    "        \"api_type\": \"openai\",  # OpenAI-compatible for function calling\n",
    "    },\n",
    "    temperature=0.3,\n",
    "    max_tokens=800,\n",
    "    # NOTE: Do NOT use tools=[...] here; pass functions to the agent instead\n",
    ")\n",
    "# Create a function-calling assistant agent\n",
    "function_assistant = AssistantAgent(\n",
    "    name=\"grok_function_assistant\",\n",
    "    system_message=(\n",
    "        \"You are a helpful assistant that can call functions to get weather information and perform calculations. \"\n",
    "        \"Use the available tools when appropriate.\"\n",
    "    ),\n",
    "    llm_config=function_config,\n",
    "    functions=[get_weather, calculate_math],  # Pass functions here\n",
    ")\n",
    "# Test function calling\n",
    "user_proxy.initiate_chat(\n",
    "    function_assistant,\n",
    "    message=\"What's the weather like in Tokyo? Also, can you calculate 15 * 23 + 7?\",\n",
    "    max_turns=2,\n",
    "    clear_history=True,\n",
    ")\n",
    "```\n",
    "**Best Practices:**\n",
    "- Register your Python functions with the agent using the `functions` argument.\n",
    "- Do **not** put `tools` or `functions` in `LLMConfig`—this will cause serialization errors.\n",
    "- Use `Annotated` for function argument descriptions to improve tool usability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from autogen import AssistantAgent, LLMConfig, UserProxyAgent\n",
    "\n",
    "# Example 1: Using Grok with default OpenAI client (OpenAI-compatible)\n",
    "grok4_config = LLMConfig(\n",
    "    {\n",
    "        \"model\": \"grok-4\",\n",
    "        \"api_type\": \"openai\",  # Use existing openai type only\n",
    "        \"base_url\": \"https://api.x.ai/v1\",\n",
    "        \"api_key\": os.getenv(\"XAI_API_KEY\"),\n",
    "        \"max_tokens\": 1000,\n",
    "        # Test if extra_body works for search parameters\n",
    "        \"extra_body\": {\n",
    "            \"search_enabled\": True,\n",
    "            \"real_time_data\": True,\n",
    "            \"search_parameters\": {\n",
    "                \"max_search_results\": 5,\n",
    "                \"include_citations\": True,\n",
    "                \"search_timeout\": 10,\n",
    "                \"return_citations\": True,\n",
    "            },\n",
    "        },\n",
    "    },\n",
    "    temperature=0.5,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 1: Basic Grok Conversation with Real-time Search\n",
    "\n",
    "This example demonstrates basic conversation with Grok, leveraging its real-time search capabilities through the `extra_body` parameter.\n",
    "\n",
    "**Key features:**\n",
    "- **Real-time data access**: Grok can access current information and web search\n",
    "- **Search parameters**: Configure max results, citations, and timeout settings\n",
    "- **Standard OpenAI compatibility**: Works with existing AG2 patterns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create agents\n",
    "assistant = AssistantAgent(\n",
    "    name=\"grok_assistant\",\n",
    "    system_message=\"You are a helpful AI assistant powered by Grok. You have access to real-time information and can help with various tasks.\",\n",
    "    llm_config=grok4_config,\n",
    ")\n",
    "\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"user_proxy\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    max_consecutive_auto_reply=3,\n",
    "    is_termination_msg=lambda x: x.get(\"content\", \"\").rstrip().endswith(\"TERMINATE\"),\n",
    "    code_execution_config={\n",
    "        \"work_dir\": \"coding\",\n",
    "        \"use_docker\": False,  # Set to True if you have Docker available\n",
    "    },\n",
    ")\n",
    "\n",
    "user_proxy.initiate_chat(\n",
    "    assistant,\n",
    "    message=\"What's the weather like in Tokyo? Also, can you calculate 15 * 23 + 7?\",\n",
    "    max_turns=2,\n",
    "    clear_history=True,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Result Analysis\n",
    "\n",
    "Notice that the conversation completed successfully! Grok was able to:\n",
    "- Provide today's date correctly (July 18, 2025)\n",
    "- Access real-time information through its built-in search capabilities\n",
    "- The search parameters in `extra_body` enabled enhanced search functionality\n",
    "\n",
    "The response shows that Grok's real-time search integration works seamlessly with AG2."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example 2: Function Calling with Grok\n",
    "\n",
    "This example demonstrates how to use Grok with function calling capabilities. This was challenging to get right due to AG2's specific requirements for function registration.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example 2: Using Grok with function calling capabilities\n",
    "print(\"\\n=== Grok with Function Calling ===\")\n",
    "\n",
    "from typing import Annotated\n",
    "\n",
    "from autogen.tools import tool\n",
    "\n",
    "\n",
    "@tool(description=\"Get current weather for a city\")\n",
    "def get_weather(city: Annotated[str, \"The city name\"]) -> str:\n",
    "    \"\"\"Get current weather for a city.\"\"\"\n",
    "    # This is a mock function - in reality, you'd call a weather API\n",
    "    return f\"The current weather in {city} is sunny with a temperature of 22°C.\"\n",
    "\n",
    "\n",
    "@tool(description=\"Calculate a mathematical expression\")\n",
    "def calculate_math(expression: Annotated[str, \"Mathematical expression to evaluate\"]) -> str:\n",
    "    \"\"\"Calculate a mathematical expression safely.\"\"\"\n",
    "    try:\n",
    "        # Simple evaluation - in production, use a safer math parser\n",
    "        result = eval(expression.replace(\"^\", \"**\"))\n",
    "        return f\"The result of {expression} is {result}\"\n",
    "    except Exception:\n",
    "        return f\"Could not evaluate the expression: {expression}\"\n",
    "\n",
    "\n",
    "# Configure Grok - FIXED: Remove tools from LLMConfig (causes JSON serialization error)\n",
    "function_config = LLMConfig(\n",
    "    {\n",
    "        \"model\": \"grok-4\",\n",
    "        \"api_key\": os.getenv(\"XAI_API_KEY\"),\n",
    "        \"base_url\": \"https://api.x.ai/v1\",\n",
    "        \"api_type\": \"openai\",  # Using OpenAI-compatible client for function calling\n",
    "    },\n",
    "    temperature=0.3,  # Temperature goes at top level, not in config_list\n",
    "    max_tokens=800,  # max_tokens also at top level\n",
    "    # NOTE: DO NOT use tools=[functions] here - causes \"Object of type function is not JSON serializable\"\n",
    ")\n",
    "\n",
    "# Create function-calling assistant - FIXED: Pass functions to ConversableAgent, not LLMConfig\n",
    "function_assistant = AssistantAgent(\n",
    "    name=\"grok_function_assistant\",\n",
    "    system_message=\"You are a helpful assistant that can call functions to get weather information and perform calculations. Use the available tools when appropriate.\",\n",
    "    llm_config=function_config,\n",
    "    functions=[get_weather, calculate_math],  # CORRECT: Functions go here, not in LLMConfig.tools\n",
    ")\n",
    "\n",
    "result = function_assistant.run(\n",
    "    message=\"What's the weather like in Tokyo? Also, can you calculate 15 * 23 + 7?\",\n",
    "    max_turns=2,\n",
    "    tools=[get_weather, calculate_math],\n",
    ")\n",
    "\n",
    "result.process()"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Grok example with live search",
   "tags": [
    "grok",
    "openai"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
